import gradio as gr
import numpy as np
import joblib
import cv2

# 加载模型
with open('best_model.pkl', 'rb') as f:
    model = joblib.load(f)  # 使用joblib.load代替pickle.load

# 修改预测函数
def predict_image(image):
    # Gradio 传入的 image 已经是 numpy 数组，直接转换为灰度图
    if len(image.shape) == 3:  # 如果是彩色图像
        img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:  # 如果已经是灰度图
        img = image
        
    # 调整尺寸为32x32
    img = cv2.resize(img, (32, 32))
    
    # 扁平化处理，适应模型输入
    image_flattened = img.flatten().reshape(1, -1)
    
    # 模型预测
    prediction = model.predict(image_flattened)[0]
    return {
        "猫": float(prediction == 0),
        "狗": float(prediction == 1)
    }

# 创建Gradio接口
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(),  # 移除 type="file"
    outputs=gr.Label(num_top_classes=2)
)

# 启动应用
iface.launch()